import random
import os
from glob import glob
import importlib
from src import const
import argparse
import pandas as pd
import numpy as np
import torch
from torch import nn

def merge_const(module_name):
    new_conf = importlib.import_module(module_name)
    for key, value in new_conf.__dict__.items():
        if not(key.startswith('_')):
            # const.__dict__[key] = value
            setattr(const, key, value)
            print('override', key, value)


def parse_args_and_merge_const():
    parser = argparse.ArgumentParser()
    parser.add_argument('--conf', default='', type=str)
    args = parser.parse_args()
    if args.conf != '':
        merge_const(args.conf)


class Evaluator(object):

    def __init__(self, label_name):
        self.label_name = label_name
        self.reset()

    def reset(self):
        # 按样本算
        self.attr_true = np.zeros(40,)
        self.attr_num = 0
        # 按样本 * 属性算
        self.attr_all_true = 0
        self.attr_all_num = 0
        self.pred_prob = []
        self.pred_label = []

    def attr_count(self, sample, output):
        pred_prob = nn.functional.softmax(output, dim=1)[:, 1, :].detach().cpu().numpy()
        pred_label = torch.argmax(output, dim=1).detach().cpu().numpy()
        self.pred_prob.append(pred_prob)
        self.pred_label.append(pred_label)
        target = sample['label'].detach().cpu().numpy()
        self.attr_num += pred_label.shape[0]
        self.attr_true += (pred_label == target).sum(axis=0)
        self.attr_all_num += pred_label.shape[0] * pred_label.shape[1]
        self.attr_all_true += (pred_label == target).sum()

    def add(self, sample, output):
        self.attr_count(sample, output)

    def evaluate(self):
        ret = {}
        attr_acc = self.attr_true / self.attr_num
        for i in range(len(self.label_name)):
            ret['acc_{}'.format(self.label_name[i])] = attr_acc[i]
        ret['acc'] = self.attr_all_true / self.attr_all_num
        return ret

    def get_pred_df(self):
        return pd.DataFrame(
            np.concatenate(self.pred_prob),
            columns=['pred_' + x for x in self.label_name]
        )
